package com.example.toy.collector.beam;

import com.example.toy.collector.RpcSender;
import com.example.toy.collector.model.RpcSenderFactory;
import com.example.toy.core.record.ToyRecord;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.beam.sdk.transforms.DoFn;

import java.io.IOException;

/**
 * BeamCollectionSender
 *
 * @author shenbing
 * @date 2021/12/17
 */
@Slf4j
@Getter
public class BeamCollectionSender extends DoFn<ToyRecord, Void> {

  RpcSender sender;

  @StartBundle
  public void startBundle() {
    sender = RpcSenderFactory.build();
  }

  @ProcessElement
  public void processElement(@Element ToyRecord element, ProcessContext context) {
    sender.send(element);
  }

  @FinishBundle
  public void finishBundle() {
    try {
      sender.close();
    } catch (IOException e) {
      log.warn("Error to close toy record sender", e);
    }
  }
}
